import tensorflow as tf
import numpy as np
import cv2
import os
from tensorflow import gfile
import matplotlib.pyplot as plt

class_names = os.listdir("../../../../../large_data/CV3/_many_files/agriculture/train")
class_names = sorted(class_names)
CHECKPOINT_DIR = "./_save/checkpoints"
test_dir = "../../../../../large_data/CV3/_many_files/agriculture/test/agriculture"
model_dir = "../../../../../large_data/model/inceptionV3/tensorflow_inception_graph.pb"
n_class = 5


def load_google_model(model_dir):
    with gfile.FastGFile(model_dir, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        tf.import_graph_def(graph_def, name="")


def create_test_featrue(test_dir):
    with tf.Session() as sess:
        st = sess.graph.get_tensor_by_name("pool_3/_reshape:0")
        test_data, test_feature, test_labels = [], [], []
        for i in os.listdir(test_dir):
            img = cv2.imread(os.path.join(test_dir, i))
            img = cv2.resize(img, (256, 256))
            img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
            test_data.append(img)
            img_data = gfile.FastGFile(os.path.join(test_dir, i), "rb").read()
            feature = sess.run(st, feed_dict={"DecodeJpeg/contents:0": img_data})
            test_feature.append(feature)
            test_labels.append(i.split("_")[0])
    return test_data, np.reshape(test_feature, (-1, 2048)), np.array(test_labels)


def show_img(test_data, pre_labels, test_labels):
    _, axs = plt.subplots(4, 4)
    for i, axi in enumerate(axs.flat):
        axi.imshow(test_data[i])
        print(pre_labels[i], test_labels[i])
        axi.set_xlabel(xlabel=pre_labels[i], color="black" if pre_labels[i] == test_labels[i] else "red")
        axi.set(xticks=[], yticks=[])
    plt.savefig(os.path.join("_save", 'agriculture' + ".jpg"))
    plt.show()


if __name__ == '__main__':
    load_google_model(model_dir)
    test_data, test_feature, test_labels = create_test_featrue(test_dir)
    # graph
    x_transfer = tf.placeholder(tf.float32, [None, 2048])
    y_transfer = tf.placeholder(tf.int64, [None])  # [None,5]
    logits = tf.layers.dense(x_transfer, n_class)
    saver = tf.train.Saver()
    with tf.Session() as sess:
        sess.run(tf.global_variables_initializer())
        print(CHECKPOINT_DIR)
        last_point = tf.train.latest_checkpoint(CHECKPOINT_DIR)
        print(last_point)
        saver.restore(sess, last_point)
        pred = sess.run(tf.argmax(logits, 1), {x_transfer: test_feature})
        show_img(test_data, [class_names[i] for i in pred], test_labels)
